# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: donggu

import datetime

from numpy.testing import *

from coin.exchange.okex.kr_rest.product import OkexProduct
import experimental.prophet as prophet
from experimental.prophet.graph import graph
from experimental.prophet.ops import timeseries

start_time = datetime.date(2019, 7, 1)
end_time = datetime.datetime(2019, 7, 1, 0, 1)


def model():
  ts = prophet.global_timestamp()
  _timer = prophet.timer('1s')

  product = OkexProduct.FromStr('BTC-USDT')
  feed = prophet.fastfeed_coin(product, 'Spot.Okex.v3')

  p = prophet.last_valid(prophet.trade(feed).price)

  p = p.named_ref('trade_price')
  p_2s = prophet.time_shift(p, '2s').named_ref('trade_price_2s')
  diff_2s = prophet.indicator.diff(p, time_window='2s').named_ref('diff_2s')
  ret_2s = prophet.time_shift_y(diff_2s, '2s').named_ref('ret_2s')
  # ratio = prophet.indicator.diff_ratio(p, time_window='20ms')
  # log_return = prophet.indicator.log_return(p, time_window='20ms')

  with prophet.control_if(ts % 10**9 == 0):
    aggregator = prophet.scalar_value_aggregator([ts, p, p_2s, diff_2s, ret_2s])

  return aggregator


def test_diff():
  df = prophet.run_from_fastfeed_multiprocess(model,
                                              start_time,
                                              end_time,
                                              machine='feed-01.ap-northeast-1.aws',
                                              use_run_cache=False,
                                              eval_updated_only=False)

  df['manual_diff_2s'] = df['trade_price'] - df['trade_price_2s']

  assert_array_equal(df['manual_diff_2s'], df['diff_2s'])
  assert_array_equal(df['diff_2s'][2:], df['ret_2s'][:-2])
